#include "G4SystemOfUnits.hh"
#include "Randomize.hh"
#include "G4ParticleTable.hh"
#include "G4IonTable.hh"
#include "G4ParticleDefinition.hh"
#include "G4PrimaryParticle.hh"
#include "G4PrimaryVertex.hh"
#include "G4HadronElastic.hh"
#include "G4Alpha.hh"
#include "G4Pow.hh"
#include <memory>

#include "MyGenCreator.hh"
#include "MyGenContainer.hh"

#include "NuclearRecoil.hh"
#include "Migdal.hh"
#include <string>

G4String remove_figure(G4String str)
{
    std::string sstr = str.data();
    int length = sstr.size();
    for (int i = 0; i < length; ++i)
    {
        if (std::isdigit(sstr[i]))
        {
            sstr.erase(i,1);
            --i;
        }
    }
    return G4String(sstr.c_str());
}

Migdal::Migdal(): MyBaseGen()
{
    m_PGType = "Nuclear Recoil Background";
    m_ClassName = "Migdal";
    m_TypeFlag = 1;
    m_EvtFlag = 0;

    m_Particle = "C";
    m_ParticleEnergy = 1. * MeV;
    m_ParticlePosition = G4ThreeVector(0., 0., 0.);
    m_ParticlePolarization = G4ThreeVector(0., 0., 0.);
    m_ParticleMomentumDirection = G4ThreeVector(0., 0., 0.);

    m_PosGen = nullptr;
    m_IonPD = nullptr;
    m_NeutronPD = nullptr;
}

Migdal::~Migdal()
{
    if(!m_FileEkin)
        delete m_FileEkin;
    if(!m_FilePDF)
        delete m_FilePDF;
}

void Migdal::Initialize(YAML::Node yamlNode)
{
    m_PosGen = MyGenContainer::GetInstance().GetMyPos();

    G4ParticleTable *particleTable = G4ParticleTable::GetParticleTable();
    m_ElectronPD = particleTable->FindParticle("e-");
    m_NeutronPD = particleTable->FindParticle("neutron");

    m_Z = yamlNode["Z"].as<int>(6); // 原子序数
    m_A = yamlNode["A"].as<double>(12); // 质量数
    m_Charge = yamlNode["Charge"].as<double>(-1e6); // 电荷
    if(m_Charge <= -1e6)
    {
        m_Charge = m_Z;
    }

    m_FileEkin = new TFile(yamlNode["HistFile"].as<std::string>("NeutronKineticEnergy.root").c_str());
    m_HistEkin = dynamic_cast<TH1D*>(m_FileEkin->Get(yamlNode["HistName"].as<std::string>("NeutronEkin").c_str()));
    m_FilePDF  = new TFile(yamlNode["ElecPDF"].as<std::string>("MigdalEnergyPDF.root").c_str());
    m_LowElecE = yamlNode["LowElectronEnergy"].as<double>(-1.);
    if (m_LowElecE < -1.)
        m_LowElecE = -1.;
}

void Migdal::GeneratePrimaryVertex(G4Event* anEvent)
{
    if(m_IonPD == nullptr)
    {
        G4IonTable *ionTable = (G4IonTable *)(G4ParticleTable::GetParticleTable()->GetIonTable());
        m_IonPD = ionTable->GetIon(m_Z, m_A);
        m_Particle = ionTable->GetIonName(m_Z, m_A);
        if(m_Z == 2 && m_A == 4)
        {
            m_IonPD = G4Alpha::Alpha();
        }
        m_FilePDF->Print();
        m_HistPDF = dynamic_cast<TH2F*>(m_FilePDF->Get(remove_figure(m_Particle)));
        m_PjIon = m_HistPDF->ProjectionY("HistPjIon");
        m_NumSlice = m_HistPDF->GetNbinsY();
        for(G4int iy = 0; iy < m_NumSlice; ++iy)
        {
            if (iy >= MAXIONHIST)
            {
                std::cerr << "Error! Number of SliceIon reach the maxlength of TH1D array!" << std::endl;
            }
            m_SliceIon[iy] = m_HistPDF->ProjectionX(TString::Format("SliceIon%d", iy), iy+1, iy+1);
        }
        std::cout << __LINE__ << std::endl;
    }

    PosAndMom neutronPosAndMomDir = m_PosGen->GenPosAndMomDir();
    G4double neutronEk = m_HistEkin->GetRandom() * MeV;
    G4DynamicParticle neutronTrack(m_NeutronPD, neutronPosAndMomDir.second, neutronEk);
    G4DynamicParticle ionTrack = NuclearRecoil::HadronNuclearScattering(neutronTrack, m_Z, m_A);
    m_ParticleEnergy = ionTrack.GetKineticEnergy();
    m_ParticlePosition = neutronPosAndMomDir.first;
    m_ParticleMomentumDirection = ionTrack.GetMomentumDirection();

    // Set Primary Vertex of ion
    G4PrimaryParticle* priIon = new G4PrimaryParticle(m_IonPD);
    priIon->SetKineticEnergy(m_ParticleEnergy);
    priIon->SetMomentumDirection(m_ParticleMomentumDirection);
    priIon->SetCharge(m_Charge);
    G4PrimaryVertex* vecIon = new G4PrimaryVertex(neutronPosAndMomDir.first, 0.);
    vecIon->SetPrimary(priIon);
    anEvent->AddPrimaryVertex(vecIon);

    // Set Primary Vertex of electron
    G4PrimaryParticle* priElc = new G4PrimaryParticle(m_ElectronPD);
    priElc->SetKineticEnergy(this->ElectronEnergy(m_ParticleEnergy));
    priElc->SetMomentumDirection(MyBaseGen::MomDir4Pi());
    G4PrimaryVertex* vecElc = new G4PrimaryVertex(neutronPosAndMomDir.first, 0.);
    vecElc->SetPrimary(priElc);
    anEvent->AddPrimaryVertex(vecElc);

    ++m_EvtFlag;
}

G4double Migdal::ElectronEnergy(G4double ionEnergy)
{
    G4int nBinIon = m_PjIon->FindBin(ionEnergy);
    if(nBinIon < 1)
        nBinIon = 1;
    if(nBinIon >= m_NumSlice)
        nBinIon = m_NumSlice;
    G4double elecEnergy = -1.e6;
    while (elecEnergy < m_LowElecE)
    {
        elecEnergy = m_SliceIon[nBinIon]->GetRandom();
    } 
    return elecEnergy * keV;
}

REGISTGPSGEN(Migdal)